import os

import mmcv
# using a pre-trained detector. 预训练配置
from mmcv import Config

from mmdet.apis import set_random_seed, inference_detector, show_result_pyplot


# root_path=r"/project/train/src_repo/mmdetection/" # 要改的

root_path=r"/home/deepin/Documents/ji_pingtai/mmdetection/" # 要改的

cfg_path=root_path+'configs/swin/retinanet_swin-t-p4-w7_fpn_3x_coco.py'

cfg = Config.fromfile(cfg_path)


#训练----------------------------------------------------
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector


cfg.model.bbox_head.num_classes = 2 #----类别数

cfg.work_dir = 'work_dir'

# Set seed thus the results are more reproducible #播种，结果更具可重复性
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)


cfg.device='cuda' # 'ConfigDict' object has no attribute 'device'


cfg.log_config.interval = 20 #打印log 间隔
cfg.evaluation.interval = 20 # 评估间隔，以减少评估时间
cfg.checkpoint_config.interval = 100  # 保存间隔:  设置检查点：，以降低存储成本
cfg.runner = dict(type='EpochBasedRunner', max_epochs=100)# ---训练 轮次




# Build dataset 构建 数据集
datasets = [build_dataset(cfg.data.train)]

# Build the detector 构建 识别
model = build_detector(cfg.model)

# Add an attribute for visualization convenience 添加属性以方便可视化，模型的类别
model.CLASSES = datasets[0].CLASSES

# Create work_dir 文件夹
mmcv.mkdir_or_exist(os.path.abspath(cfg.work_dir))
#保存 cfg 配置：
# 把自定义的 config 在 自定义的working dir 下 保存一份------------推理和测试 可以用到的
cfg.dump(F'{cfg.work_dir}/swin_3.py')


train_detector(model, datasets, cfg, distributed=False, validate=True) # -------进行----训练-----


# # # 查看 评估结果 load tensorboard in colab
# # %load_ext tensorboard
# #
# # # see curves in tensorboard
# # tensorboard --logdir ./tutorial_exps  #--------运行这个看看 数的 曲线



